'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: keywordTableService.py
* @Time: 2025/11/13
* @All Rights Reserve By Brtc
'''
from dataclasses import dataclass
from uuid import UUID

from injector import inject
from redis import Redis

from pkg.sqlalchemy import SQLAlchemy
from .base_service import BaseService
from internal.model import KeywordTable, Segment
from ..entity.cache_entity import LOCK_KEYWORD_TABLE_UPDATE_KEYWORD_TABLE, LOCK_EXPIRE_TIME


@inject
@dataclass
class KeywordTableService(BaseService):
    """知识库关键词检索服务"""
    db:SQLAlchemy
    redis_client:Redis


    def get_keyword_table_from_dataset_id(self, dataset_id:UUID)->KeywordTable:
        """根据传递的知识库id获取关键词表,如果没有则新建一个新的"""
        keyword_table = self.db.session.query(KeywordTable).filter(
            KeywordTable.dataset_id == dataset_id).one_or_none()
        if keyword_table is None:
            keyword_table = self.create(KeywordTable,
                                        dataset_id=dataset_id,
                                        keyword_table={})
        return keyword_table

    def add_keyword_table_from_ids(self, dataset_id:UUID, segment_ids: list[UUID])->None:
        """根据传递的知识库id  +  片段 id 列表在关键词表中添加关键词"""
        #1、新增加关键词， 该操作需要上锁， 避免 出错
        cache_key = LOCK_KEYWORD_TABLE_UPDATE_KEYWORD_TABLE.format(dataset_id=dataset_id)
        with self.redis_client.lock(cache_key, LOCK_EXPIRE_TIME):
            #2、获取指定知识库的关键词表
            keyword_table_record = self.get_keyword_table_from_dataset_id(dataset_id)
            keyword_table = {
                field:set(value) for field, value in keyword_table_record.keyword_table.items()
            }
            #3、根据segment_ids 查找片段的 关键词信息
            segments = self.db.session.query(Segment).with_entities(Segment.id, Segment.keywords).filter(
                Segment.id.in_(segment_ids),
            ).all()

            #4、循环将新关键词添加到关键词表中
            for id, keywords in segments:
                for keyword in keywords:
                    if keyword not in keyword_table:
                        keyword_table[keyword] = set()
                    keyword_table[keyword].add(str(id))
            #5、更新关键词表
            self.update(
                keyword_table_record,
                keyword_table={field:list(value) for field, value in keyword_table.items()},
            )


    def delete_keyword_table_from_ids(self, dataset_id:UUID, segment_ids:list[UUID])->None:
        """根据传递的知识库id  +  片段 id 列表在关键词表中删除关键词"""
        # 1、新增加关键词， 该操作需要上锁， 避免 出错
        cache_key = LOCK_KEYWORD_TABLE_UPDATE_KEYWORD_TABLE.format(dataset_id=dataset_id)
        with self.redis_client.lock(cache_key, LOCK_EXPIRE_TIME):
            #2、获取当前知识库的关键词表
            keyword_table_record = self.get_keyword_table_from_dataset_id(dataset_id)
            keyword_table = keyword_table_record.keyword_table.copy()
            #3、将片段id 列表转换成集合, 并创建关键词用于清空关键词
            segment_ids_to_delete = set([str(segment_id) for segment_id in segment_ids])
            keyword_to_delete = set()
            #4、循环遍历所有关键词执行判断与更新
            for keyword, ids in keyword_table.items():
                ids_set = set(ids)
                if segment_ids_to_delete.intersection(ids_set):
                    keyword_table[keyword] = list(ids_set.difference(segment_ids_to_delete))
                    if not keyword_table[keyword]:
                        keyword_to_delete.add(keyword)
            #5、检测空关键词数据并删除（关键词并没有映射任何字段的id 数据）
            for keyword in keyword_to_delete:
                del keyword_table[keyword]
            #6、将数据更新都关键此表中
            self.update(keyword_table_record, keyword_table=keyword_table)















